from ultralytics import YOLO
from src.config import STAGE1_CONFIG
import shutil
import os


def run_training_stage1():
    print("--- Starting Stage 1 Training: YOLO Proposer ---")

    if os.path.exists(STAGE1_CONFIG['weights_path']):
        print(f"Existing weights detected: {STAGE1_CONFIG['weights_path']}, will fine-tune from this checkpoint.")
        model_path = STAGE1_CONFIG['weights_path']
    else:
        print(f"No local weights found, loading pretrained '{STAGE1_CONFIG['model_name']}' model from ultralytics.")
        model_path = STAGE1_CONFIG['model_name']

    model = YOLO(model_path)

    try:
        results = model.train(
            data=STAGE1_CONFIG['data_yaml'],
            epochs=STAGE1_CONFIG['epochs'],
            batch=STAGE1_CONFIG['batch_size'],
            imgsz=STAGE1_CONFIG['img_size'],
            project=os.path.join("runs", "train"),
            name="stage1_proposer",
            exist_ok=True,  # Allow overwriting previous training results
            save=True,  # Ensure checkpoints and final model are saved
            verbose=True  # Display detailed training logs
        )
    except Exception as e:
        print(f"\n[ERROR] An error occurred during YOLO training: {e}")
        print("Please check:")
        print(f"  1. Dataset paths are correctly configured in '{STAGE1_CONFIG['data_yaml']}'.")
        print("  2. Dataset files are complete and valid.")
        print("  3. The `ultralytics` library and all dependencies are installed.")
        return

    source_best_weights = os.path.join("runs", "train", "stage1_proposer", "weights", "best.pt")

    if os.path.exists(source_best_weights):
        os.makedirs(os.path.dirname(STAGE1_CONFIG['weights_path']), exist_ok=True)
        shutil.copyfile(source_best_weights, STAGE1_CONFIG['weights_path'])
        print(f"\nStage 1 training complete! Best model copied from '{source_best_weights}' to '{STAGE1_CONFIG['weights_path']}'")

        # Update SERVER_CONFIG for subsequent use
        from src.config import SERVER_CONFIG
        SERVER_CONFIG['stage1_weights'] = STAGE1_CONFIG['weights_path']
    else:
        print(
            f"\n[ERROR] Training appears complete, but the best model was not found at '{source_best_weights}'. Please check training logs to diagnose the issue.")


if __name__ == '__main__':
    if not os.path.exists(STAGE1_CONFIG['data_yaml']):
        print(f"Error: Dataset configuration file not found at '{STAGE1_CONFIG['data_yaml']}'.")
        print("Please create and configure your data.yaml file in the 'dataset/' directory.")
    else:
        run_training_stage1()

